Conversation
|
Per MAINTAINER.md, CODEOWNERS has requested the Speculative decoding owners. Flagging the KV Cache oncalls as well — @ispobock @xiezhq-hermann — since the change is in Two points that may help review:
Non-hybrid-SWA models are unaffected: the whole body is behind |
spec_prepare_for_decode routes dflash-family algorithms (DFLASH, DSPARK) to their own prepare_for_decode, which never calls maybe_evict_swa and never advances decode_batch_idx; eagle_prepare_for_decode does both. On hybrid-SWA models this retains sliding-window KV for every generated token, so a single request exhausts the SWA pool after swa_full_tokens_ratio * max_total_num_tokens generated tokens and is retracted by the scheduler. Pool sizing already assumes the eviction runs (pool_configurator: trailing_tokens = sliding_window + eviction_interval * draft_tokens + page_size). Mirror eagle_prepare_for_decode: evict before advancing decode_batch_idx so the overlap-scheduler gate (decode_batch_idx >= 1) keeps protecting the first decode iteration. Reproduced on DeepSeek-V4-Flash-0731 (SM120, TP=2, DSPARK block 5, ratio 0.1, pool 778,496): SWA occupancy tracked generated tokens 1:1 and the request was retracted at exactly 77,824 tokens. With the fix, the same request generated 100,000 tokens with SWA occupancy holding at ~512 tokens.
26a2a39 to
40034c4
Compare
…SH decode path `spec_prepare_for_decode` routes the dflash family (which includes DSPARK) to `batch.spec_info.prepare_for_decode`, and that handler omits two things the EAGLE path does: `maybe_evict_swa()` and advancing `decode_batch_idx`. Both are needed -- the eviction gate inside `maybe_evict_swa` is `swa_maintenance_step and req.decode_batch_idx >= 1`, so adding the call alone would still never fire. The effect on a hybrid-SWA model is that SWA KV is retained for every generated token and never reclaimed: measured on DeepSeek-V4-Flash-0731 (tp8 + CP8 + DSPARK, A100), a single request at 33,792 tokens of context held 32,300 SWA tokens instead of a bounded window. The pool then fills at `swa_full_tokens_ratio * max_total_num_tokens` and the scheduler starts retracting, which re-prefills `prompt + everything generated so far`. Pool sizing already assumes this eviction runs (pool_configurator computes `trailing_tokens = window + eviction_interval * draft_tokens + page_size`). Measured at in=1024 / out=32768, ratio 0.1 unless noted, no retraction in any row: before conc=11 1113.6 tok/s before conc=14, ratio 0.3 1331.9 tok/s swa peak 0.91 <- best reachable after conc=32 2220.1 tok/s swa peak 0.14 after conc=48 2292.8 tok/s swa peak 0.20 1.72x, and per-request SWA at full context drops 32,300 -> 1,024. Because SWA stops being the scarce pool, `--swa-full-tokens-ratio` no longer needs raising: 0.1 measures the same as 0.3 (2220 vs 2186 tok/s, inside run-to-run noise), so the default is right again and the memory goes back to the full pool. Upstream PR is still open; applies clean on this branch. (cherry picked from commit 3226bc5882ecaefa29d52081137197ff33c1d436)
3 participants
Motivation
spec_prepare_for_decoderoutes dflash-family algorithms (DFLASH, DSPARK) toDFlashDraftInputV2.prepare_for_decode, which never callsmaybe_evict_swaand never advancesdecode_batch_idx— which the eviction gate requires to be>= 1.eagle_prepare_for_decodedoes both. On hybrid-SWA models the sliding-window eviction is therefore unreachable under DFLASH/DSPARK: SWA KV is retained for every generated token, so a single request exhausts the SWA pool after roughlyswa_full_tokens_ratio * max_total_num_tokensgenerated tokens and is retracted by the scheduler mid-generation.Reproduced on DeepSeek-V4-Flash-0731 (2x RTX PRO 6000, SM120, TP=2, DSPARK block size 5,
swa_full_tokens_ratio=0.1,max_total_num_tokens=778496).#swa tokentracks generated tokens 1:1, and the request is retracted at exactly 77,824 tokens —align_page_size(int(778496 * 0.1)), the full SWA pool — while the full pool is only 10% used:Long prompts are unaffected — a 131k-token prefill holds ~8.5K SWA tokens. Only generated tokens accumulate.
The same missing call also skips the DSV4 compress-state drain (
maybe_evict_dsv4_state, documented "called every decode step from ScheduleBatch"), so the NPU + DSV4 + DFLASH combination leaks state-pool slots for the same reason.Modifications
speculative/spec_utils.py: in the dflash-family branch, mirroreagle_prepare_for_decode— callbatch.maybe_evict_swa()and advancedecode_batch_idx. Eviction runs before the tick so the overlap-scheduler gate (decode_batch_idx >= 1) still skips the first decode iteration, exactly as on the eagle path.test/registered/unit/spec/test_decode_bookkeeping_ownership.py: record the new bookkeeping owner in_OWNER_SITES, as that guard requires for a genuinely new owner.test/registered/unit/speculative/test_spec_prepare_swa_eviction.py(new): asserts the dflash branch evicts, ticks, evicts before ticking, keeps advancing across iterations (so the counter stays a clock rather than a flag), and that the eagle path is unchanged.Accuracy Tests
Eviction frees only SWA slots behind the window —
max(sliding_window, page_size)back from the sequence on the radix path, page-flooredpre_len - sliding_windowon the chunk-cache path. SWA layers attend within the window, and both the eagle and non-spec paths already call this same code every decode step (eagle_utils.py,mem_cache/allocation.py). No output change observed.Speed Tests and Profiling
The
ignore_eosrequest that previously died at 77,809 tokens completed 100,000 tokens with the fix, at a sustained ~159 tok/s and zero retractions.#swa tokensampled repeatedly across the run stayed between 512 and 1,024 (~1% of the SWA pool), against 77,824 at the point of failure before the fix.Checklist
Prepared with AI assistance.
CI States
Latest PR Test (Base): ❌ Run #31830154702
Latest PR Test (Extra): ❌ Run #31830154341